package com.feidee.fd.sml.algorithm.feature

import com.feidee.fd.sml.algorithm.component.feature.{PolynomialExpander, PolynomialExpanderParam}
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.apache.spark.ml.linalg.Vectors
import org.scalatest.FunSuite

/**
  * @Author songhaicheng
  * @Date 2019/3/20 14:00
  * @Description
  * @Reviewer
  */
class PolynomialExpanderSuite extends FunSuite {

  val paramStr: String =
    """
      |{
      |	"inputCol": "col0",
      | "outputCol": "features",
      |	"preserveCols": "col1,col2"
      |}
    """.stripMargin

  val pe = new PolynomialExpander()
  val param: PolynomialExpanderParam = pe.parseParam(new ToolClass().encrypt(paramStr))

  /**
    * PolynomialExpander 参数构造方法测试：检测到读取出合法参数 & 成功赋值默认参数，即认为测试通过
    */
  test("PolynomialExpander parameter construction test") {
    assert(param.degree == 2 &&
      "col0".equals(param.inputCol) &&
      "features".equals(param.outputCol) &&
      "col1,col2".equals(param.preserveCols))
  }

  /**
    * PolynomialExpander 模型训练功能测试：传入生成的测试数据和正确参数，可以正常训练 PolynomialExpander 模型，
    * 且预测结果列等于设置的列（outputCol + preserveCols），即认为测试通过
    */
  test("PolynomialExpander training function test") {
    val data = Seq(
      (Vectors.dense(2.0, 1.0),  "col10", "col20", "col30"),
      (Vectors.dense(0.0, 0.0),  "col11", "col21", "col31"),
      (Vectors.dense(3.0, -1.0),  "col11", "col22", "col32")
    )
    val df = TestingDataGenerator.spark.createDataFrame(data).toDF("col0", "col1", "col2", "col3")

    val pe = new PolynomialExpander()
    val polyDF = pe.train(param, df).transform(df)
    polyDF.show(truncate = false)
    assert(polyDF.schema.size == 3)
  }

}
